import librosa
import numpy as np
import scipy.io.wavfile
import os
import matplotlib.pyplot as plt
import seaborn as sns
import scipy.signal

VER = 'v1.0'
BASE_DIR, FILE_NAME = os.path.split(__file__)
path = '../../../../../large_data/audio/zsn-stop1.wav'
AUDIO_PATH = os.path.join(BASE_DIR, path)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)

signal01, sr01 = librosa.load(AUDIO_PATH, sr=None, res_type='kaiser_fast')
print('sr01', sr01)
print('signal01', signal01.shape)
sr02, signal02 = scipy.io.wavfile.read(AUDIO_PATH)
signal02 = signal02[:, 0]
print('sr02', sr02)
print('signal02', signal02.shape)

spr = 2
spc = 1
spn = 0
plt.figure(figsize=[8, 8])
spn += 1
plt.subplot(spr, spc, spn)
plt.title('from librosa')
plt.plot(signal01)
spn += 1
plt.subplot(spr, spc, spn)
plt.title('from scipy.io.wavfile')
plt.plot(signal01)
plt.show()

mfcc = librosa.feature.mfcc(signal01, sr=sr01, n_mfcc=50)
print('mfcc', mfcc.shape)
nperseg = int(sr02 / 1000. * 25)
noverlap = nperseg - int(sr02 / 1000. * 10)
print('nperseg', nperseg)
print('noverlap', noverlap)
_, _, spectrum = scipy.signal.spectrogram(
    signal02,
    fs=sr02,
    window='hann',
    nperseg=nperseg,
    noverlap=noverlap,
    detrend=False,
)
print('spectrum', spectrum.shape)
spectrum = np.float32(spectrum)
spectrum = np.where(spectrum <= 0., np.finfo(np.float32).eps, spectrum)
spectrum = np.log10(spectrum)
print('spectrum', spectrum.shape)

spr = 1
spc = 3
spn = 0
plt.figure(figsize=[12, 6])
spn += 1
plt.subplot(spr, spc, spn)
plt.title('mfcc seaborn')
sns.heatmap(mfcc)
spn += 1
plt.subplot(spr, spc, spn)
plt.title('mfcc plt.imshow')
plt.imshow(mfcc)
spn += 1
plt.subplot(spr, spc, spn)
plt.title('scipy.signal.spectrogram')
sns.heatmap(spectrum)
plt.show()
